'''
runProblem.py - Paul Talbot, NE 553, Oregon State University, Spring 2011
Used in conjunction with buildMesh.py and phiSolver.py.

This code is the handler for the classes created in buildMesh and phiSolver.
It writes output files, spawns neutrons to traverse the mesh, stores K and F, and
checks tolerance to convergence.
'''
import buildMesh #includes all mesh-related classes
import phiSolver as ps #includes phi-converging functions
import time

F=[1,0,0] #array of F-values current, older, oldest
K=[1,0,0] #array of K-values current, older, oldest
tolK=10**-2 #tolerances for convergence
tolF=10**-2
tolPhi=10**-3

def solvePhi(mesh): #iterates on scalar flux
  pc.doPhi(mesh,tolPhi)

def findAverage(list):
  total=0
  for item in list:
    total+=item
  return total/len(list)

def runProblem(case=0,quad=0,mesh=0):
  initTime=time.clock()
  if mesh==0: #no mesh given
    myMesh=buildMesh.Mesh(case,quad,tolPhi) #initialize the mesh
  else: myMesh=mesh #mesh supplied
  #define output files
  fileOut=file('out.'+myMesh.case,'w') #writes K, F each generation
  filePhi=file('phi.'+myMesh.case,'w') #writes phi for each cell, fast and thermal
  filePsi=file('psi.'+myMesh.case,'w') #writes psi (fast, thermal)(minus, average, plus)
  #fileCur=file('cur.'+myMesh.case,'w') #writes current in each cell (fast,thermal)(L, Total, R)
  #put headers in output files
  fileOut.writelines('gen ; F ; K ; Time ; Generations ; Fast Iter. ; Therm. Iter.\n')
  filePhi.writelines('Region ; Cell ; Position ; Fast Flux ; Thermal Flux\n')
  #fileCur.writelines('Region ; Cell ; Position ; Fast L ; Fast R ; Fast T ; ;'+\
  #    ' Thermal L ; Thermal R ; Thermal T\n')
  filePsi.writelines('Region ; Cell ; Position ; ')
  for n in myMesh.mus:
    filePsi.writelines('Angle '+str(n)[:5]+' ; Fast - ; Fast i ; Fast + ;'+\
        'Thermal - ; Thermal i ; Thermal + ; ;')
  filePsi.writelines('\n')
  
  converged=False #start out unconverged
  gen=0
  while not converged: #keep going until F is converged
    starttime=time.time()
    gen+=1 #next iteration started
    F[1:]=F[:2] #store old values
    K[1:]=K[:2]
    F[0]=K[0]=0 #reset upcoming value

    #converge phi values
    i=ps.doPhi(myMesh,tolPhi) #does fast(R,L), converge, thermal(R,L), converge
    
    myMesh.getFdx() #calculate F on mesh
    F[0]=myMesh.Fdx #store it
    checkF=ps.testTolerance(F,tolF) #check tolerance

    myMesh.getK() #calculate K on mesh
    K[0]=myMesh.K #store it
    checkK=ps.testTolerance(K,tolK) #check tolerance

    if gen>1:
      if checkK and checkF:converged=True #both tolerated, then trigger converged

    #display to screen
    elapsed=time.time()-starttime
    if gen%1==0:
      print gen,'F:',str(F[0])[:6],'K:',str(K[0])[:6],\
          'Time:',str(elapsed)[:6],'sec','Fast Iter.:',i[0],\
          'Thermal Iter.:',i[1] #debug or watch-while-run
    #writing outputs
    fileOut.writelines(str(gen)+';'+str(F[0])+';'+str(K[0])+';'+\
        str(elapsed)+';'+str(gen)+';'+str(i[0])+';'+str(i[1])+'\n') #write K, F to file
    fileOut.flush() #make sure it writes, not buffers
    #if gen>5:converged=True #hack for debugging

  #after convergence, write most recent data
  for rgn in myMesh.rgns:
    for cell in rgn.cells:
      #write data common to all outputs
      for f in [filePhi,filePsi]: #fileCur left out
        f.writelines(str(cell.rgn.num)+';'+str(cell.num)+';'+str(cell.pos)+';')
      #write phi data
      filePhi.writelines(str(cell.g[0].phi[0])+';'+str(cell.g[1].phi[0])+'\n') 
      #write current data
      #fileCur.writelines(str(cell.g[0].curL)+';'+str(cell.g[0].curR)+';'+\str(cell.g[0].curT)+';'+\
      #    ' ;'+str(cell.g[1].curL)+';'+str(cell.g[1].curR)+';'+\str(cell.g[1].curT)+'\n')
      #write psi data
      for n in range(len(myMesh.mus)):
        filePsi.writelines(' ; '+\
            str(+cell.g[0].psim[n])+';'+str(cell.g[0].psia[n])+';'\
                +str(cell.g[0].psip[n])+';'+str(cell.g[1].psim[n])+\
                ';'+str(cell.g[1].psia[n])+';'+str(cell.g[1].psip[n])+'; ;')
      filePsi.writelines('\n')
  #write fuel pin average data
  filePin=file('pin.'+myMesh.case,'w')
  filePin.writelines('Pin ; Av. Fast Flux ; Av. Thermal Flux\n')
  for pin in myMesh.pins:
    pin.getPhi(1)
    pin.getPhi(2)
    filePin.writelines(str(pin.num)+';'+str(pin.groups[0].phi)+';'+str(pin.groups[1].phi)+'\n')
  #write assembly average data
  fileAbly=file('assembly.'+myMesh.case,'w')
  fileAbly.writelines('Material ; Av. Fast Flux ; Av. Thermal Flux ; Fast fL ; Fast fR ; '+\
      'Thermal fL ; Thermal fR ; Fast Tot Xsec ; Fast Abs Xsec ; Fast Fis Xsec ; Fast sIn Xsec ;'+\
      ' Fast sOut Xsec ; Thermal Tot Xsec ; Thermal Abs Xsec ; Thermal Fis Xsec ; '+\
      'Thermal sIn Xsec ; Thermal sOut Xsec\n')
  print myMesh
  for g in [1,2]:
    myMesh.groups[g-1].avPhi()
    myMesh.groups[g-1].homogenize()
  fileAbly.writelines(myMesh.case+';'+str(myMesh.groups[0].phi)+';'+str(myMesh.groups[1].phi)+';'+\
      str(myMesh.groups[0].fL)+';'+str(myMesh.groups[0].fR)+';'+str(myMesh.groups[1].fL)+';'+\
      str(myMesh.groups[1].fR)+';'+\
      str(myMesh.groups[0].sigt)+';'+str(myMesh.groups[0].siga)+';'+\
      str(myMesh.groups[0].sigf)+';'+str(myMesh.groups[0].sigs_in)+';'+\
      str(myMesh.groups[0].sigs_out)+';'+str(myMesh.groups[1].sigt)+';'+\
      str(myMesh.groups[1].siga)+';'+\
      str(myMesh.groups[1].sigf)+';'+str(myMesh.groups[1].sigs_in)+';'+\
      str(myMesh.groups[1].sigs_out)+'\n')

  elapsedTime=time.clock()-initTime
  print '\nTotal run time for',myMesh.case,':',elapsedTime,'\n'
if __name__=='__main__':
  runProblem(0)
